{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71f30092-a437-48cb-beb0-e820a4d7deb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://zhuanlan.zhihu.com/p/578610797"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc418a3-4882-4d9e-abe1-5a1fc291aeae",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.distributed as dist\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from torch.nn.parallel import DistributedDataParallel\n",
    "\n",
    "def setup():\n",
    "    dist.init_process_group('nccl')\n",
    "\n",
    "\n",
    "def cleanup():\n",
    "    dist.destroy_process_group()\n",
    "\n",
    "\n",
    "class ToyModel(nn.Module):\n",
    "\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__()\n",
    "        self.layer = nn.Linear(1, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layer(x)\n",
    "\n",
    "\n",
    "class MyDataset(Dataset):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.data = torch.tensor([1, 2, 3, 4], dtype=torch.float32)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.data[index:index + 1]\n",
    "\n",
    "\n",
    "ckpt_path = 'tmp.pth'\n",
    "\n",
    "\n",
    "def main():\n",
    "    setup()\n",
    "    rank = dist.get_rank()\n",
    "    pid = os.getpid()\n",
    "    print(f'current pid: {pid}')\n",
    "    print(f'Current rank {rank}')\n",
    "    device_id = rank % torch.cuda.device_count()\n",
    "\n",
    "    dataset = MyDataset()\n",
    "    sampler = DistributedSampler(dataset)\n",
    "    dataloader = DataLoader(dataset, batch_size=2, sampler=sampler)\n",
    "\n",
    "    model = ToyModel().to(device_id)\n",
    "    ddp_model = DistributedDataParallel(model, device_ids=[device_id])\n",
    "    loss_fn = nn.MSELoss()\n",
    "    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)\n",
    "\n",
    "    if rank == 0:\n",
    "        torch.save(ddp_model.state_dict(), ckpt_path)\n",
    "\n",
    "    dist.barrier()\n",
    "\n",
    "    map_location = {'cuda:0': f'cuda:{device_id}'}\n",
    "    state_dict = torch.load(ckpt_path, map_location=map_location)\n",
    "    print(f'rank {rank}: {state_dict}')\n",
    "    ddp_model.load_state_dict(state_dict)\n",
    "\n",
    "    for epoch in range(2):\n",
    "        sampler.set_epoch(epoch)\n",
    "        for x in dataloader:\n",
    "            print(f'epoch {epoch}, rank {rank} data: {x}')\n",
    "            x = x.to(device_id)\n",
    "            y = ddp_model(x)\n",
    "            optimizer.zero_grad()\n",
    "            loss = loss_fn(x, y)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "    cleanup()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
